package decision_tree.dao;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import dao.FileIO;
import decision_tree.UnigramModel;
import decision_tree.dto.TreeNode;
import decision_tree.question.Question;

public class DecisionTreeIO {
	protected QuestionIO questionIO = new QuestionIO();
	public void dump(TreeNode root, String filename) {
		Map<TreeNode, Integer> nodeMap = new HashMap<TreeNode, Integer>();
		getMap(root, nodeMap, 0);
		
		List<String> list = new LinkedList<String>();
		list.add("NUM_STATES\t" + nodeMap.size() + "\r\n");
		list.add("ROOT_STATES\t" + nodeMap.get(root) + "\r\n");
		list.add("***Question\r\n");
		list.addAll(getQuestion(nodeMap));
		list.add("***Unigram\r\n");
		list.addAll(getUnigramModel(nodeMap));
		list.add("***Inheritance\r\n");
		list.addAll(getStructure(nodeMap, root));
		
		FileIO.writeFile(filename, list);
	}
	
	public TreeNode load(String filename) {
		List<String> list = FileIO.readFile(filename);
		int nStates = Integer.parseInt(list.get(0).split("\t")[1]);
		int rState = Integer.parseInt(list.get(1).split("\t")[1]);
		
		List<String> questionList = getList(list, "***Question");
		List<String> unigramList = getList(list, "***Unigram");
		List<String> inheritanceList = getList(list, "***Inheritance");
		
		Map<Integer, TreeNode> nodeMap = getNodeMap(nStates);
		
		Map<TreeNode, Question> questionMap = getQuestion(questionList, nodeMap);
		Map<TreeNode, UnigramModel> unigramMap = getUnigram(unigramList, nodeMap);
		Map<TreeNode, Map<Integer, TreeNode>> inheritanceMap = getInheritance(inheritanceList, nodeMap);
		
		integrate(questionMap, unigramMap, inheritanceMap);
	
		return nodeMap.get(rState);
	}
	
	// dump *********************************
	
	protected int getMap(TreeNode root, Map<TreeNode, Integer> map, int index) {
		map.put(root, index);
		index ++;
		if(root.children != null){
			for(TreeNode child : root.children.values()) {
				index = getMap(child, map, index);
			}
		}
		return index;
	}
	
	protected List<String> getQuestion(Map<TreeNode, Integer> nodeMap) {
		List<String> list = new ArrayList<String>();
		for(Entry<TreeNode, Integer> entry : nodeMap.entrySet()) {
			Question q = entry.getKey().q;
			if(q != null) {
				String question = questionIO.write(q);
				list.add(entry.getValue() + "\t" + question + "\r\n");
			}
		}
		return list;
	}
	
	protected List<String> getUnigramModel(Map<TreeNode, Integer> nodeMap) {
		List<String> list = new ArrayList<String> ();
		for(Entry<TreeNode, Integer> entry : nodeMap.entrySet()) {
			UnigramModel um = entry.getKey().model;
			if(um != null) {
				int index = entry.getValue();
				String model = um.toLine();
				list.add(index + "\t" + model + "\r\n");
			}
		}
		return list;
	}
	
	protected List<String> getStructure(Map<TreeNode, Integer> nodeMap, TreeNode root) {
		List<Inheritance> inheritanceList = getInheritances(nodeMap);
		List<String> list = new ArrayList<String>(inheritanceList.size());
		for(Inheritance inheritance : inheritanceList) {
			int from = nodeMap.get(inheritance.from);
			int to = nodeMap.get(inheritance.to);
			int label = inheritance.label;
			
			String str = from + "\t" + to + "\t" + label + "\r\n";
			list.add(str);
		}
		return list;
	}
	
	protected List<Inheritance> getInheritances(Map<TreeNode, Integer> nodeMap) {
		List<Inheritance> list = new ArrayList<Inheritance>();
		for(Entry<TreeNode, Integer> entry : nodeMap.entrySet()) {
			TreeNode fromNode = entry.getKey();
			if(fromNode.children != null) {
				for(Entry<Integer, TreeNode> subEntry : fromNode.children.entrySet()) {
					Inheritance inheritance = new Inheritance();
					inheritance.from = fromNode;
					inheritance.to = subEntry.getValue();
					inheritance.label = subEntry.getKey();
					
					list.add(inheritance);
				}
			}
		}
		return list;
	}
	
	// load *********************************
	
	protected void integrate(Map<TreeNode, Question> questionMap, Map<TreeNode, UnigramModel> unigramMap, Map<TreeNode, Map<Integer, TreeNode>> inheritanceMap) {
		// add question
		for(Entry<TreeNode, Question> entry : questionMap.entrySet()) {
			entry.getKey().q = entry.getValue();
		}
		
		// add unigram
		for(Entry<TreeNode, UnigramModel> entry : unigramMap.entrySet()) {
			entry.getKey().model = entry.getValue();
		}
		
		// add inheritance
		for(Entry<TreeNode, Map<Integer, TreeNode>> entry : inheritanceMap.entrySet()) {
			entry.getKey().children = entry.getValue();
		}
	}
	
	protected Map<Integer, TreeNode> getNodeMap(int n) {
		Map<Integer, TreeNode> map = new HashMap<Integer, TreeNode>();
		for(int i = 0; i < n; i++) {
			map.put(i, new TreeNode());
		}
		return map;
	}
	
	protected Map<TreeNode, Question> getQuestion(List<String> list, Map<Integer, TreeNode> nodeMap) {
		Map<Integer, Question> map = getQuestion(list);
		Map<TreeNode, Question> result = new HashMap<TreeNode, Question>();
		for(Entry<Integer, Question> entry : map.entrySet()) {
			result.put(nodeMap.get(entry.getKey()), entry.getValue());
		}
		return result;
	}
	
	protected Map<Integer, Question> getQuestion(List<String> list) {
		Map<Integer, Question> map = new HashMap<Integer, Question>();
		for(String line : list) {
			String [] parts = line.split("\t");
			int index = Integer.parseInt(parts[0]);
			Question q = questionIO.parse(parts[1]);
			map.put(index, q);
		}
		return map;
	}
	
	public Map<TreeNode, UnigramModel> getUnigram(List<String> list, Map<Integer, TreeNode> nodeMap) {
		Map<Integer, UnigramModel> map = getUnigram(list);
		Map<TreeNode, UnigramModel> result = new HashMap<TreeNode, UnigramModel>();
		for(Entry<Integer, UnigramModel> entry : map.entrySet()) {
			result.put(nodeMap.get(entry.getKey()), entry.getValue());
		}
		return result;
	}
	
	public Map<Integer, UnigramModel> getUnigram(List<String> list) {
		Map<Integer, UnigramModel> map = new HashMap<Integer, UnigramModel>();
		for(String line : list) {
			String [] parts = line.split("\t");
			int index = Integer.parseInt(parts[0]);
			UnigramModel model = UnigramModel.parse(parts[1]);
			map.put(index, model);
		}
		return map;
	}
	
	protected Map<TreeNode, Map<Integer, TreeNode>> getInheritance(List<String> list, Map<Integer, TreeNode> nodeMap) {
		Map<Integer, Map<Integer, Integer>> map = getInheritance(list);
		Map<TreeNode, Map<Integer, TreeNode>> result = new HashMap<TreeNode, Map<Integer, TreeNode>>();
		for(Entry<Integer, Map<Integer, Integer>> entry : map.entrySet()) {

			Map<Integer, Integer> subMap = entry.getValue();
			Map<Integer, TreeNode> subResult = new HashMap<Integer, TreeNode>();
			for(Entry<Integer, Integer> subEntry : subMap.entrySet()) {
				int label = subEntry.getKey();
				TreeNode toNode = nodeMap.get(subEntry.getValue());
				subResult.put(label, toNode);
			}
			
			TreeNode fromNode = nodeMap.get(entry.getKey());
			result.put(fromNode, subResult);
		}
		return result;
	}
	
	protected Map<Integer, Map<Integer, Integer>> getInheritance(List<String> list) {
		List<List<Integer>> flat = loadInheritance(list);
		Map<Integer, Map<Integer, Integer>> map = new HashMap<Integer, Map<Integer, Integer>>();
		
		for(List<Integer> entry : flat) {
			int from = entry.get(0);
			int to = entry.get(1);
			int label = entry.get(2);
			
			if(!map.containsKey(from)) {
				map.put(from, new HashMap<Integer, Integer>());
			}
			
			map.get(from).put(label, to);
		}
		return map;
	}
	
	protected List<List<Integer>> loadInheritance(List<String> list) {
		List<List<Integer>> result = new ArrayList<List<Integer>>(list.size());
		for(String line : list) {
			String [] parts = line.split("\t");
			int from = Integer.parseInt(parts[0]);
			int to = Integer.parseInt(parts[1]);
			int label = Integer.parseInt(parts[2]);
			result.add(Arrays.asList(from, to, label));
		}
		return result;
	}
	
	protected List<String> getList(List<String> list, String key) {		
		Iterator<String> itr = list.iterator();
		int index = 0;
		boolean started = false;
		int from = 0;
		int to = list.size();
		
		while(itr.hasNext()) {
			String line = itr.next();
			if(line.startsWith("***")) {
				if(!started) {
					if(line.compareTo(key) == 0) {
						from = index + 1;
						started = true;
					}
				} else {
					to = index;
					break;
				}
			}
			index ++;
		}
		return list.subList(from, to);
	}
	
	// structure **************************************
	
	private static class Inheritance {
		public TreeNode from;
		public TreeNode to;
		public int label;
	}
}
